import argparse
import os
import xml.etree.ElementTree as ET

from utils.basic import get_file_name
from utils.basic import merge_dict


def parse_args():
    parser = argparse.ArgumentParser(
        description='Prepare lists txt file for dataset')
    parser.add_argument(
        '--dataset',
        dest='dataset',
        help='dataset to use',
        default='data',
        type=str)
    parser.add_argument(
        '--target',
        dest='target',
        help='output list file',
        default='data',
        type=str)
    parser.add_argument(
        '--root',
        dest='root_path',
        help='dataset root path',
        default='',
        type=str)
    parser.add_argument(
        '--postfix',
        dest='postfix',
        help='postfix to file',
        default='',
        type=str)
    parser.add_argument(
        '--prefix',
        dest='prefix',
        help='prefix to file',
        default='',
        type=str)
    parser.add_argument(
        '--label_file',
        dest='label_file',
        help='label file',
        default='',
        type=str)
    parser.add_argument(
        '--class_num',
        dest='class_num',
        help='class number',
        default=9,
        type=int)
    parser.add_argument(
        '--instance',
        dest='instance',
        help='required instance',
        default=0,
        type=int)
    return parser.parse_args()


def check_xml_and_img_file(file_path, lists, params):
    """
    a filter for traverse file system
    :param file_path: file absolution path
    :param lists: a series parameters needed, it likes:
                   0: file list of all the required file,
                   1: dict of the class name as the key and the number of
                   the class as the value
    :param params: some additional parameters, it likes:
                    0: class name list, the required class names
                    1: a string for split the different files
    :return: lists
    """
    try:
        tree = ET.parse(file_path)
        root = tree.getroot()
        size = root.find('size')
        if not ((float(size.find('width').text) > 0) and
                (float(size.find('height').text) > 0)):
            return lists
        class_num_dict = {}
        for obj in root.iter('object'):
            string = obj.find('name').text
            if len(params[0]) > 0:
                if string not in params[0]:
                    continue
            if string not in class_num_dict:
                class_num_dict[string] = 1
            else:
                class_num_dict[string] += 1
        if len(class_num_dict) > 0:
            image_path = os.path.join(os.path.dirname(file_path),
                                      get_file_name(file_path))
            if os.path.exists(image_path + '.jpg'):
                img_path = image_path + '.jpg'
            elif os.path.exists(image_path + '.png'):
                img_path = image_path + '.png'
            elif os.path.exists(image_path + '.JPG'):
                img_path = image_path + '.JPG'
            elif os.path.exists(image_path + '.PNG'):
                img_path = image_path + '.PNG'
            else:
                return lists
            print(img_path)
            class_num_dict = merge_dict(class_num_dict, lists[1])
            lists[0].append(img_path + params[1] + file_path)
            lists[1] = class_num_dict
        return lists
    except ET.ParseError:
        return lists
